import argparse
import multiprocessing
import sys
import signal
import os
from tcp_forwarder import tcp_forward
from udp_forwarder import udp_forward
from security import validate_port, validate_host
from logger import setup_logger
from config import load_config, get_default_forwarding, get_forwarding_rules, set_config_file

# 设置日志记录器
logger = setup_logger("main")

def parse_arguments():
    """解析命令行参数"""
    # 获取默认配置
    default_config = get_default_forwarding()
    
    parser = argparse.ArgumentParser(description='端口转发工具 - 支持 TCP 和 UDP')
    parser.add_argument('-l', '--local-port', type=int, 
                        default=default_config.get("local_port"),
                        help=f'本地监听端口 (默认: {default_config.get("local_port", 8080)})')
    parser.add_argument('-r', '--remote-host', type=str, 
                        default=default_config.get("remote_host"),
                        help=f'远程主机地址 (默认: {default_config.get("remote_host", "localhost")})')
    parser.add_argument('-p', '--remote-port', type=int, 
                        default=default_config.get("remote_port"),
                        help=f'远程主机端口 (默认: {default_config.get("remote_port", 80)})')
    parser.add_argument('-t', '--protocol', type=str, 
                        choices=['tcp', 'udp', 'both'], 
                        default=default_config.get("protocol", "tcp"), 
                        help=f'协议类型: tcp, udp 或 both (默认: {default_config.get("protocol", "tcp")})')
    parser.add_argument('-c', '--config-rule', type=str,
                        help='使用配置文件中的规则名称')
    parser.add_argument('-a', '--all-rules', action='store_true',
                        help='启动配置文件中所有启用的规则')
    parser.add_argument('--config', type=str,
                        help='指定配置文件路径')
    
    args = parser.parse_args()
    
    # 如果指定了配置文件路径
    if args.config:
        if os.path.exists(args.config):
            set_config_file(args.config)
            logger.info(f"使用配置文件: {args.config}")
        else:
            logger.error(f"指定的配置文件不存在: {args.config}")
            sys.exit(1)
    
    # 如果指定了使用配置规则
    if args.config_rule:
        rules = get_forwarding_rules()
        for rule in rules:
            if rule.get("name") == args.config_rule:
                args.local_port = rule.get("local_port")
                args.remote_host = rule.get("remote_host")
                args.remote_port = rule.get("remote_port")
                args.protocol = rule.get("protocol")
                break
        else:
            logger.error(f"找不到名为 '{args.config_rule}' 的规则")
            sys.exit(1)
    
    return args

def setup_signal_handlers(tcp_pool=None, udp_monitor=None):
    """设置信号处理器以确保资源清理"""
    def signal_handler(sig, frame):
        logger.info("正在关闭端口转发...")
        if tcp_pool:
            tcp_pool.cleanup()
        if udp_monitor:
            udp_monitor.stop()
        logger.info("端口转发已停止")
        sys.exit(0)
        
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)

def start_forwarding(local_port, remote_host, remote_port, protocol):
    """启动单个转发规则"""
    # 添加参数验证
    if not validate_port(local_port):
        logger.error(f"错误: 无效的本地端口 {local_port}")
        return False
        
    if not validate_host(remote_host):
        logger.error(f"错误: 无效的远程主机 {remote_host}")
        return False
        
    if not validate_port(remote_port):
        logger.error(f"错误: 无效的远程端口 {remote_port}")
        return False
    
    logger.info(f"启动端口转发: {protocol.upper()} 模式, 本地端口: {local_port}, 远程: {remote_host}:{remote_port}")
    
    try:
        if protocol.lower() == 'tcp':
            # 创建连接池
            from connection_pool import ConnectionPool
            tcp_pool = ConnectionPool(remote_host, remote_port)
            
            # 设置信号处理
            setup_signal_handlers(tcp_pool=tcp_pool)
            
            tcp_forward(local_port, remote_host, remote_port)
            return True
        elif protocol.lower() == 'udp':
            udp_forward(local_port, remote_host, remote_port)
            return True
        else:  # both
            logger.info("同时启动 TCP 和 UDP 转发")
            # 创建进程池
            with multiprocessing.Pool(processes=2) as pool:
                # 异步执行 TCP 和 UDP 转发
                tcp_result = pool.apply_async(tcp_forward, (local_port, remote_host, remote_port))
                udp_result = pool.apply_async(udp_forward, (local_port, remote_host, remote_port))
                
                # 保持主进程运行
                try:
                    # 等待进程结束（实际上不会结束，除非出错或被中断）
                    tcp_result.get()
                    udp_result.get()
                except KeyboardInterrupt:
                    logger.info("端口转发已停止")
                    pool.terminate()
            return True
                    
    except KeyboardInterrupt:
        logger.info("端口转发已停止")
        return False
    except Exception as e:
        logger.error(f"端口转发出错: {e}", exc_info=True)
        return False

def main():
    """主函数"""
    # 输出配置文件路径，方便调试
    from config import CONFIG_FILE
    
    args = parse_arguments()
    
    logger.info(f"配置文件路径: {CONFIG_FILE}")
    
    # 如果指定了启动所有规则
    if args.all_rules:
        rules = get_forwarding_rules()
        if not rules:
            logger.warning("配置文件中没有启用的转发规则")
            return
        
        # 创建进程列表
        processes = []
        
        # 为每个规则创建一个进程
        for rule in rules:
            p = multiprocessing.Process(
                target=start_forwarding,
                args=(
                    rule.get("local_port"),
                    rule.get("remote_host"),
                    rule.get("remote_port"),
                    rule.get("protocol", "tcp")
                )
            )
            processes.append(p)
            p.start()
            logger.info(f"已启动规则: {rule.get('name')}")
        
        # 等待所有进程完成
        try:
            for p in processes:
                p.join()
        except KeyboardInterrupt:
            logger.info("正在停止所有转发...")
            for p in processes:
                if p.is_alive():
                    p.terminate()
            logger.info("所有转发已停止")
    else:
        # 启动单个转发
        start_forwarding(args.local_port, args.remote_host, args.remote_port, args.protocol)

if __name__ == "__main__":
    main()